{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "cc10e0e5-7e11-4a84-b8fd-a47be71ad185",
   "metadata": {},
   "source": [
    "# Transforms Evaluation\n",
    "\n",
    "<a href=\"https://colab.research.google.com/github/run-llama/llama_index/blob/main/docs/examples/transforms/TransformsEval.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>\n",
    "\n",
    "Here we try out different transformations and evaluate their quality.\n",
    "- First we try out different parsers (PDF, JSON)\n",
    "- Then we try out different extractors"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c756ee66",
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install llama-index-readers-file\n",
    "%pip install llama-index-llms-openai\n",
    "%pip install llama-index-embeddings-openai"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c1fcef2e-b8ef-4858-a441-8655489dc340",
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install llama-index"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2d5943b1-5e20-4067-903b-cda622ca8921",
   "metadata": {},
   "source": [
    "## Load Data + Setup\n",
    "\n",
    "Load in the Tesla data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "98423735-11af-4e28-85aa-db3d150e53d0",
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "\n",
    "pd.set_option(\"display.max_rows\", None)\n",
    "pd.set_option(\"display.max_columns\", None)\n",
    "pd.set_option(\"display.width\", None)\n",
    "pd.set_option(\"display.max_colwidth\", None)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "798eae7b-9985-4eab-a39d-ea431f4e9179",
   "metadata": {},
   "outputs": [],
   "source": [
    "!wget \"https://www.dropbox.com/scl/fi/mlaymdy1ni1ovyeykhhuk/tesla_2021_10k.htm?rlkey=qf9k4zn0ejrbm716j0gg7r802&dl=1\" -O tesla_2021_10k.htm\n",
    "!wget \"https://www.dropbox.com/scl/fi/rkw0u959yb4w8vlzz76sa/tesla_2020_10k.htm?rlkey=tfkdshswpoupav5tqigwz1mp7&dl=1\" -O tesla_2020_10k.htm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8ecf1d0a-2d58-4457-959d-696c817831b7",
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.readers.file import FlatReader\n",
    "from pathlib import Path\n",
    "\n",
    "reader = FlatReader()\n",
    "docs = reader.load_data(Path(\"./tesla_2020_10k.htm\"))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b5083792-0879-44bc-8596-131c3e8560f3",
   "metadata": {},
   "source": [
    "## Generate Eval Dataset / Define Eval Functions\n",
    "\n",
    "Generate a \"golden\" eval dataset from the Tesla documents.\n",
    "\n",
    "Also define eval functions for running a pipeline."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "de8da417-04d2-4d87-b3f3-41bfef2e45aa",
   "metadata": {},
   "source": [
    "Here we define an ingestion pipeline purely for generating a synthetic eval dataset."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a2d673b4-8226-4612-ac48-338e3cb44fd0",
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.core.evaluation import DatasetGenerator, QueryResponseDataset\n",
    "from llama_index.llms.openai import OpenAI\n",
    "from llama_index.embeddings.openai import OpenAIEmbedding\n",
    "from llama_index.readers.file import FlatReader\n",
    "from llama_index.core.node_parser import HTMLNodeParser, SentenceSplitter\n",
    "from llama_index.core.ingestion import IngestionPipeline\n",
    "from pathlib import Path\n",
    "\n",
    "import nest_asyncio\n",
    "\n",
    "nest_asyncio.apply()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "424cd4e5-2fab-4769-bdf4-b148d9b77b20",
   "metadata": {},
   "outputs": [],
   "source": [
    "reader = FlatReader()\n",
    "docs = reader.load_data(Path(\"./tesla_2020_10k.htm\"))\n",
    "\n",
    "pipeline = IngestionPipeline(\n",
    "    documents=docs,\n",
    "    transformations=[\n",
    "        HTMLNodeParser.from_defaults(),\n",
    "        SentenceSplitter(chunk_size=1024, chunk_overlap=200),\n",
    "        OpenAIEmbedding(),\n",
    "    ],\n",
    ")\n",
    "eval_nodes = pipeline.run(documents=docs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ee14925f-45ca-4780-b457-b30f03c933fd",
   "metadata": {},
   "outputs": [],
   "source": [
    "# NOTE: run this if the dataset isn't already saved\n",
    "# Note: we only generate from the first 20 nodes, since the rest are references\n",
    "# eval_llm = OpenAI(model=\"gpt-4-1106-preview\")\n",
    "eval_llm = OpenAI(model=\"gpt-3.5-turbo\")\n",
    "\n",
    "dataset_generator = DatasetGenerator(\n",
    "    eval_nodes[:100],\n",
    "    llm=eval_llm,\n",
    "    show_progress=True,\n",
    "    num_questions_per_chunk=3,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7982e1e1-f72a-422f-ab76-f8a5b59b6d54",
   "metadata": {},
   "outputs": [],
   "source": [
    "eval_dataset = await dataset_generator.agenerate_dataset_from_nodes(num=100)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cd54a416-eb5c-4c4b-91d0-561db57e5ba6",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "100"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(eval_dataset.qr_pairs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fd94c286-8088-48dd-b014-6a0d125a1faa",
   "metadata": {},
   "outputs": [],
   "source": [
    "eval_dataset.save_json(\"data/tesla10k_eval_dataset.json\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "06a224bd-f919-45b4-b0f1-8c542cab8ab7",
   "metadata": {},
   "outputs": [],
   "source": [
    "# optional\n",
    "eval_dataset = QueryResponseDataset.from_json(\n",
    "    \"data/tesla10k_eval_dataset.json\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9510bdf0-4d37-41c8-8ce5-a23160221dae",
   "metadata": {},
   "outputs": [],
   "source": [
    "eval_qs = eval_dataset.questions\n",
    "qr_pairs = eval_dataset.qr_pairs\n",
    "ref_response_strs = [r for (_, r) in qr_pairs]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "97f8e3a7-4903-4717-944e-7995d953bdf1",
   "metadata": {},
   "source": [
    "### Run Evals"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fe1ef466-2520-4347-9170-06e18be468b9",
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.core.evaluation import (\n",
    "    CorrectnessEvaluator,\n",
    "    SemanticSimilarityEvaluator,\n",
    ")\n",
    "from llama_index.core.evaluation.eval_utils import (\n",
    "    get_responses,\n",
    "    get_results_df,\n",
    ")\n",
    "from llama_index.core.evaluation import BatchEvalRunner"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1352239c-6f2b-4b77-a8e6-ddd831af7db3",
   "metadata": {},
   "outputs": [],
   "source": [
    "evaluator_c = CorrectnessEvaluator(llm=eval_llm)\n",
    "evaluator_s = SemanticSimilarityEvaluator(llm=eval_llm)\n",
    "evaluator_dict = {\n",
    "    \"correctness\": evaluator_c,\n",
    "    \"semantic_similarity\": evaluator_s,\n",
    "}\n",
    "batch_eval_runner = BatchEvalRunner(\n",
    "    evaluator_dict, workers=2, show_progress=True\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e6e4ee4f-2ca2-455f-9449-54a7c30d224b",
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.core import VectorStoreIndex\n",
    "\n",
    "\n",
    "async def run_evals(\n",
    "    pipeline, batch_eval_runner, docs, eval_qs, eval_responses_ref\n",
    "):\n",
    "    # get query engine\n",
    "    nodes = pipeline.run(documents=docs)\n",
    "    # define vector index (top-k = 2)\n",
    "    vector_index = VectorStoreIndex(nodes)\n",
    "    query_engine = vector_index.as_query_engine()\n",
    "\n",
    "    pred_responses = get_responses(eval_qs, query_engine, show_progress=True)\n",
    "    eval_results = await batch_eval_runner.aevaluate_responses(\n",
    "        eval_qs, responses=pred_responses, reference=eval_responses_ref\n",
    "    )\n",
    "    return eval_results"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3925c3d5-cadc-425e-958c-ea1ba585ff2d",
   "metadata": {},
   "source": [
    "## 1. Try out Different Sentence Splitter (Overlaps)\n",
    "\n",
    "The chunking strategy matters! Here we try the sentence splitter with different overlap values, to see how it impacts performance.\n",
    "\n",
    "The `IngestionPipeline` lets us concisely define an e2e transformation pipeline for RAG, and we define variants where each corresponds to a different sentence splitter configuration (while keeping other steps fixed)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "963551fe-eb1f-4254-a533-1ffd01cbc364",
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.core.node_parser import HTMLNodeParser, SentenceSplitter\n",
    "\n",
    "# For clarity in the demo, make small splits without overlap\n",
    "sent_parser_o0 = SentenceSplitter(chunk_size=1024, chunk_overlap=0)\n",
    "sent_parser_o200 = SentenceSplitter(chunk_size=1024, chunk_overlap=200)\n",
    "sent_parser_o500 = SentenceSplitter(chunk_size=1024, chunk_overlap=600)\n",
    "\n",
    "html_parser = HTMLNodeParser.from_defaults()\n",
    "\n",
    "parser_dict = {\n",
    "    \"sent_parser_o0\": sent_parser_o0,\n",
    "    \"sent_parser_o200\": sent_parser_o200,\n",
    "    \"sent_parser_o500\": sent_parser_o500,\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b5e64615-eda1-46b9-98a6-0e0f995b85dc",
   "metadata": {},
   "source": [
    "Define a separate pipeline for each parser."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ebef94ea-626d-428a-96fc-9d5f1e6adfd2",
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.embeddings.openai import OpenAIEmbedding\n",
    "from llama_index.core.ingestion import IngestionPipeline\n",
    "\n",
    "# generate a pipeline for each parser\n",
    "# keep embedding model fixed\n",
    "pipeline_dict = {}\n",
    "for k, parser in parser_dict.items():\n",
    "    pipeline = IngestionPipeline(\n",
    "        documents=docs,\n",
    "        transformations=[\n",
    "            html_parser,\n",
    "            parser,\n",
    "            OpenAIEmbedding(),\n",
    "        ],\n",
    "    )\n",
    "    pipeline_dict[k] = pipeline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d5cf8eb7-fec8-4c0b-87fb-5335f167bfb3",
   "metadata": {},
   "outputs": [],
   "source": [
    "eval_results_dict = {}\n",
    "for k, pipeline in pipeline_dict.items():\n",
    "    eval_results = await run_evals(\n",
    "        pipeline, batch_eval_runner, docs, eval_qs, ref_response_strs\n",
    "    )\n",
    "    eval_results_dict[k] = eval_results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "30619839-1769-403c-80e4-4b81ac0d8591",
   "metadata": {},
   "outputs": [],
   "source": [
    "# [tmp] save eval results\n",
    "import pickle\n",
    "\n",
    "pickle.dump(eval_results_dict, open(\"eval_results_1.pkl\", \"wb\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "410a02af-4d25-496d-940a-52910afb8e5d",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>names</th>\n",
       "      <th>correctness</th>\n",
       "      <th>semantic_similarity</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>sent_parser_o0</td>\n",
       "      <td>4.310</td>\n",
       "      <td>0.972838</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>sent_parser_o200</td>\n",
       "      <td>4.335</td>\n",
       "      <td>0.978842</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>sent_parser_o500</td>\n",
       "      <td>4.270</td>\n",
       "      <td>0.971759</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "              names  correctness  semantic_similarity\n",
       "0    sent_parser_o0        4.310             0.972838\n",
       "1  sent_parser_o200        4.335             0.978842\n",
       "2  sent_parser_o500        4.270             0.971759"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "eval_results_list = list(eval_results_dict.items())\n",
    "\n",
    "results_df = get_results_df(\n",
    "    [v for _, v in eval_results_list],\n",
    "    [k for k, _ in eval_results_list],\n",
    "    [\"correctness\", \"semantic_similarity\"],\n",
    ")\n",
    "display(results_df)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d5232842-1554-4350-b75a-4dbddd3110c9",
   "metadata": {},
   "outputs": [],
   "source": [
    "# [optional] persist cache in folders so we can reuse\n",
    "for k, pipeline in pipeline_dict.items():\n",
    "    pipeline.cache.persist(f\"./cache/{k}.json\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "aed975b8-a39d-48cc-b8e0-76389535a37b",
   "metadata": {},
   "source": [
    "## 2. Try out Different Extractors\n",
    "\n",
    "Similarly, metadata extraction can be quite important for good performance. We experiment with this as a last step in an overall ingestion pipeline, and define different ingestion pipeline variants corresponding to different extractors."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "092f9f69-332c-4453-b42f-9e285e3d4741",
   "metadata": {},
   "source": [
    "We define the set of document extractors we want to try out. \n",
    "\n",
    "We keep the parsers fixed (HTML parser, sentence splitter w/ overlap 200) and the embedding model fixed (OpenAIEmbedding)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5185d7f2-5af0-42a9-9159-85b01f529f77",
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.core.extractors import (\n",
    "    TitleExtractor,\n",
    "    QuestionsAnsweredExtractor,\n",
    "    SummaryExtractor,\n",
    ")\n",
    "from llama_index.core.node_parser import HTMLNodeParser, SentenceSplitter\n",
    "\n",
    "# generate a pipeline for each extractor\n",
    "# keep embedding model fixed\n",
    "extractor_dict = {\n",
    "    # \"title\": TitleExtractor(),\n",
    "    \"summary\": SummaryExtractor(in_place=False),\n",
    "    \"qa\": QuestionsAnsweredExtractor(in_place=False),\n",
    "    \"default\": None,\n",
    "}\n",
    "\n",
    "# these are the parsers that will run beforehand\n",
    "html_parser = HTMLNodeParser.from_defaults()\n",
    "sent_parser_o200 = SentenceSplitter(chunk_size=1024, chunk_overlap=200)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "05a83180-e5f6-44a5-924b-1b1b7d80f840",
   "metadata": {},
   "outputs": [],
   "source": [
    "pipeline_dict = {}\n",
    "html_parser = HTMLNodeParser.from_defaults()\n",
    "for k, extractor in extractor_dict.items():\n",
    "    if k == \"default\":\n",
    "        transformations = [\n",
    "            html_parser,\n",
    "            sent_parser_o200,\n",
    "            OpenAIEmbedding(),\n",
    "        ]\n",
    "    else:\n",
    "        transformations = [\n",
    "            html_parser,\n",
    "            sent_parser_o200,\n",
    "            extractor,\n",
    "            OpenAIEmbedding(),\n",
    "        ]\n",
    "\n",
    "    pipeline = IngestionPipeline(transformations=transformations)\n",
    "    pipeline_dict[k] = pipeline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aa704341-8ea2-479c-b521-e4a049c1eb48",
   "metadata": {},
   "outputs": [],
   "source": [
    "eval_results_dict_2 = {}\n",
    "for k, pipeline in pipeline_dict.items():\n",
    "    eval_results = await run_evals(\n",
    "        pipeline, batch_eval_runner, docs, eval_qs, ref_response_strs\n",
    "    )\n",
    "    eval_results_dict_2[k] = eval_results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "eddb8008-e0c9-4fab-9baf-97ec876e8e6a",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>names</th>\n",
       "      <th>correctness</th>\n",
       "      <th>semantic_similarity</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>summary</td>\n",
       "      <td>4.315</td>\n",
       "      <td>0.976951</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>qa</td>\n",
       "      <td>4.355</td>\n",
       "      <td>0.978807</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>default</td>\n",
       "      <td>4.305</td>\n",
       "      <td>0.978451</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "     names  correctness  semantic_similarity\n",
       "0  summary        4.315             0.976951\n",
       "1       qa        4.355             0.978807\n",
       "2  default        4.305             0.978451"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "eval_results_list_2 = list(eval_results_dict_2.items())\n",
    "\n",
    "results_df = get_results_df(\n",
    "    [v for _, v in eval_results_list_2],\n",
    "    [k for k, _ in eval_results_list_2],\n",
    "    [\"correctness\", \"semantic_similarity\"],\n",
    ")\n",
    "display(results_df)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4e8e1d27-c393-4d6f-86de-69d847a2c585",
   "metadata": {},
   "outputs": [],
   "source": [
    "# [optional] persist cache in folders so we can reuse\n",
    "for k, pipeline in pipeline_dict.items():\n",
    "    pipeline.cache.persist(f\"./cache/{k}.json\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "44157072-6b47-49dd-a0db-90e972343564",
   "metadata": {},
   "source": [
    "## 3. Try out Multiple Extractors (with Caching)\n",
    "\n",
    "TODO\n",
    "\n",
    "Each extraction step can be expensive due to LLM calls. What if we want to experiment with multiple extractors? \n",
    "\n",
    "We take advantage of **caching** so that all previous extractor calls are cached, and we only experiment with the final extractor call. The `IngestionPipeline` gives us a clean abstraction to play around with the final extractor.\n",
    "\n",
    "Try out different extractors "
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "llama_index_v2",
   "language": "python",
   "name": "llama_index_v2"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
